import os
import sys
os.system('pip install pynvml -i http://repo.myhuaweicloud.com/repository/pypi/simple --trusted-host=repo.myhuaweicloud.com')
import argparse
import moxing as mox


def download_gpu_files(dst_cache_path):
    try:
        s3_region = os.environ['S3_REGION']
    except:
        s3_region = os.environ['S3_ENDPOINT'].split('.')[1]
    if s3_region == "cn-south-1":
        s3_dir = 's3://bucket-huanan-common/moniter_tools/'
    elif s3_region == "cn-north-1":
        s3_dir = 's3://selflearning2-north1/moniter_tools/'
    elif s3_region == "cn-east-3":
        s3_dir = 's3://bucket-camera1/moniter_tools/'
    elif s3_region == "cn-north-r":
        s3_dir = 's3://bucket-bj4-sz/users/cfc/moniter_tools/'
    else:
        print('region: {} is not support yet. Please Contact wx5330805.')
        return -1 
    if not os.path.exists(dst_cache_path):
        os.makedirs(dst_cache_path)
    mox.file.copy_parallel(s3_dir, dst_cache_path)
    return 0


if __name__ == "__main__":
    """

    入参：
        boot_file=start_tele_train.py
    """

    # 获取命令行参数
    command_len = len(sys.argv)
    cmd = ''
    for i in range(1, command_len):
        print('---->', sys.argv[i])
        if 'boot_file' in sys.argv[i]:
            continue
        cmd += sys.argv[i]
        cmd += ' '

    # 获取执行路径
    exec_path = sys.argv[0]
    print('exec_path(sys.argv[0]): ', exec_path)
    start_path = exec_path[:exec_path.rfind('/')]
    exec_path = exec_path.split('/')[-2]
    print('start_path:', start_path)
    print('exec_path:', exec_path)  # code

    # 获取执行工程目录的绝对地址
    # os.path.realpath()： 当前脚本绝对路径
    # os.path.dirname()：去掉脚本的文件名
    exec_abspath = os.path.abspath(os.path.dirname(os.path.realpath(__file__)))
    print('exec_abspath:', exec_abspath)  # /home/ma-user/modelarts/user-job-dir/code

    # 获取命令行参数
    parser = argparse.ArgumentParser(description='')
    parser.add_argument('--boot_file', type=str, default="", help='boot python file')
    args, unknown = parser.parse_known_args()
    print('args:', args)  # Namespace(boot_file='start_tele_train.py')

    if args.boot_file == "":
        print('Please config boot_file!')
        sys.exit(1)

    # 从桶中下载监控文件
    ret = download_gpu_files(exec_abspath)
    if ret != 0:
        print('Cannot download moniter tools files!')
        sys.exit(1)

    # 命令组装
    cmd = '--exec_path ' + exec_path + ' ' + cmd

    bash_cmd = 'bash run_pipe_v2.sh ' + 'python ' + args.boot_file + ' ' + cmd
    print('cd %s; pwd; ls -alt; %s' % (start_path, bash_cmd))
    os.system('cd %s; pwd; ls -alt; %s' % (start_path, bash_cmd))
